import os
import time

from fuzzer.PathGreyBoxFuzzer import PathGreyBoxFuzzer
from runner.FunctionCoverageRunner import FunctionCoverageRunner
from schedule.PathPowerSchedule import PathPowerSchedule
from samples.Samples import sample1, sample2, sample3, sample4
from utils.ObjectUtils import dump_object, load_object


class Result:
    def __init__(self, coverage, crashes, start_time, end_time):
        self.covered_line = coverage
        self.crashes = crashes
        self.start_time = start_time
        self.end_time = end_time

    def __str__(self):
        return "Covered Lines: " + str(self.covered_line) + ", Crashes Num: " + str(self.crashes) + ", Start Time: " + str(self.start_time) + ", End Time: " + str(self.end_time)


if __name__ == "__main__":
    output_dir = "_result"
    os.makedirs(output_dir, exist_ok=True)

    sample_configs = [
        {"name": "Sample-1", "function": sample1, "corpus_file": "corpus/corpus_1"},
        {"name": "Sample-2", "function": sample2, "corpus_file": "corpus/corpus_2"},
        {"name": "Sample-3", "function": sample3, "corpus_file": "corpus/corpus_3"},
        {"name": "Sample-4", "function": sample4, "corpus_file": "corpus/corpus_4"},
    ]

    for config in sample_configs:
        sample_name = config["name"]
        sample_function = config["function"]
        corpus_path = config["corpus_file"]

        print("--- 开始对 " + sample_name + " 进行模糊测试 ---")

        f_runner = FunctionCoverageRunner(sample_function)
        seeds = load_object(corpus_path)

        grey_fuzzer = PathGreyBoxFuzzer(seeds=seeds, schedule=PathPowerSchedule(5), is_print=True)
        start_time = time.time()
        grey_fuzzer.runs(f_runner, run_time=600)
        res = Result(grey_fuzzer.covered_line, set(grey_fuzzer.crash_map.values()), start_time, time.time())
        dump_object(output_dir + os.sep + sample_name + "_output.pkl", res)
        print(load_object(output_dir + os.sep + sample_name + "_output.pkl"))
        print("--- " + sample_name + " 模糊测试完成 ---\n")